import pyqtgraph as pg
from PyQt6.QtWidgets import (QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, 
                           QCheckBox, QComboBox, QGroupBox, QSplitter, QDialogButtonBox)
from PyQt6.QtCore import Qt, QTimer, pyqtSignal
from utils.ohlc_data_manager import ohlc_data_manager
from utils.logger import logger
import numpy as np
from datetime import datetime
import time

# 自定义K线图元素
class CandlestickItem(pg.GraphicsObject):
    def __init__(self, data):
        pg.GraphicsObject.__init__(self)
        self.data = data  # OHLC数据
        self.picture = None
        self.generatePicture()
    
    def generatePicture(self):
        self.picture = pg.QtGui.QPicture()
        p = pg.QtGui.QPainter(self.picture)
        
        w = 0.4  # 蜡烛宽度
        for item in self.data:
            # 转换时间戳为图表x坐标
            t = item['timestamp'] / 1000  # 转为秒
            
            # 提取OHLC值
            open_price = item['open']
            high_price = item['high']
            low_price = item['low']
            close_price = item['close']
            
            # 确定颜色（上涨为红色，下跌为绿色）
            if open_price > close_price:
                color = pg.mkPen('g')
                fill = pg.mkBrush('g')
            else:
                color = pg.mkPen('r')
                fill = pg.mkBrush('r')
            
            # 绘制蜡烛线（上下影线）
            p.setPen(color)
            p.drawLine(pg.QtCore.QPointF(t, low_price), pg.QtCore.QPointF(t, high_price))
            
            # 绘制蜡烛实体
            if abs(open_price - close_price) > 0.001:  # 防止极小差异的情况
                # 绘制矩形表示开盘收盘范围
                p.setPen(color)
                p.setBrush(fill)
                rect = pg.QtCore.QRectF(
                    t - w/2,  # 左
                    open_price,  # 上
                    w,  # 宽
                    close_price - open_price  # 高
                )
                p.drawRect(rect)
            else:
                # 开盘收盘价几乎相等，绘制一条线
                p.setPen(color)
                p.drawLine(
                    pg.QtCore.QPointF(t - w/2, open_price),
                    pg.QtCore.QPointF(t + w/2, close_price)
                )
        
        p.end()
    
    def paint(self, p, *args):
        p.drawPicture(0, 0, self.picture)
    
    def boundingRect(self):
        return pg.QtCore.QRectF(self.picture.boundingRect())


class ChartViewerDialog(QDialog):
    # 自定义信号
    viewport_changed = pyqtSignal(int, tuple)  # 方向, (开始,结束)
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("K线图表查看器")
        self.setMinimumSize(1000, 600)
        self.chunk_items = {}  # 存储已加载的数据块
        self.current_view_index = -1  # 当前查看的数据块索引
        self.is_loading = False  # 防止重复加载
        self.init_ui()
        self.load_initial_data()
    
    def init_ui(self):
        layout = QVBoxLayout(self)
        
        # 图表控制区
        control_layout = QHBoxLayout()
        
        # 图表类型选择
        chart_type_group = QGroupBox("图表类型")
        chart_type_layout = QHBoxLayout()
        self.type_combo = QComboBox()
        self.type_combo.addItems(["K线图", "趋势图", "成交量"])
        self.type_combo.currentIndexChanged.connect(self.on_chart_type_changed)
        chart_type_layout.addWidget(self.type_combo)
        chart_type_group.setLayout(chart_type_layout)
        control_layout.addWidget(chart_type_group)
        
        # 时间范围控制
        time_range_group = QGroupBox("时间范围")
        time_range_layout = QHBoxLayout()
        self.zoom_combo = QComboBox()
        self.zoom_combo.addItems(["全部", "最近100条", "最近200条", "最近500条"])
        self.zoom_combo.currentIndexChanged.connect(self.on_zoom_changed)
        time_range_layout.addWidget(self.zoom_combo)
        time_range_group.setLayout(time_range_layout)
        control_layout.addWidget(time_range_group)
        
        # 数据信息显示
        self.info_label = QLabel()
        control_layout.addWidget(self.info_label)
        
        # 不再需要导航按钮
        # self.prev_btn = QPushButton("◀ 上一页")
        # self.next_btn = QPushButton("下一页 ▶")
        # self.prev_btn.clicked.connect(self.load_prev_chunk)
        # self.next_btn.clicked.connect(self.load_next_chunk)
        # control_layout.addWidget(self.prev_btn)
        # control_layout.addWidget(self.next_btn)
        
        layout.addLayout(control_layout)
        
        # 创建图表区域
        self.chart_widget = pg.PlotWidget()
        self.chart_widget.setBackground('w')  # 白色背景
        self.chart_widget.showGrid(x=True, y=True)
        
        # 添加交互特性
        self.chart_widget.setMouseEnabled(x=True, y=True)
        self.chart_widget.enableAutoRange()
        
        # 创建X轴时间格式
        axis = pg.DateAxisItem(orientation='bottom')
        self.chart_widget.setAxisItems({'bottom': axis})
        
        # 添加十字光标
        self.crosshair_v = pg.InfiniteLine(angle=90, movable=False)
        self.crosshair_h = pg.InfiniteLine(angle=0, movable=False)
        self.chart_widget.addItem(self.crosshair_v, ignoreBounds=True)
        self.chart_widget.addItem(self.crosshair_h, ignoreBounds=True)
        
        # 添加数据提示标签
        self.tooltip_label = pg.TextItem(text="", anchor=(0,1), fill=pg.mkBrush('w'), color=(0,0,0))
        self.chart_widget.addItem(self.tooltip_label)
        
        # 鼠标移动处理
        self.chart_widget.scene().sigMouseMoved.connect(self.mouse_moved)
        
        # 视图变化信号处理
        self.chart_widget.sigRangeChanged.connect(self.on_view_changed)
        
        layout.addWidget(self.chart_widget)
        
        # 底部状态栏
        status_layout = QHBoxLayout()
        self.status_label = QLabel("准备就绪")
        self.position_label = QLabel("")
        status_layout.addWidget(self.status_label)
        status_layout.addWidget(self.position_label)
        layout.addLayout(status_layout)
        
        # 绑定自定义信号
        self.viewport_changed.connect(self.handle_viewport_change)
    
    def load_initial_data(self):
        """加载全部数据"""
        data_info = ohlc_data_manager.get_data_info()
        self.info_label.setText(
            f"K线数据: {data_info['granularity']} | "
            f"数据量: {data_info['total_count']} | "
            f"时间范围: {data_info['time_range']}"
        )
        
        # 获取全部数据
        all_data = ohlc_data_manager.get_all_data()
        
        if all_data:
            self.status_label.setText(f"正在加载 {len(all_data)} 条K线数据...")
            self.load_all_data(all_data)
            self.status_label.setText(f"已加载 {len(all_data)} 条K线数据")
        else:
            self.status_label.setText("没有可用的K线数据")
    
    def create_volume_bars(self, chunk_data):
        """创建成交量柱状图"""
        timestamps = [item['timestamp']/1000 for item in chunk_data]  # 转为秒
        volumes = [item['volume'] for item in chunk_data]
        
        # 使用直方图或条形图
        # 创建红色和绿色的笔和画刷
        red_brush = pg.mkBrush('r')
        green_brush = pg.mkBrush('g')
        
        # 准备数据
        bars = []
        
        # 判断时间序列类型
        if len(timestamps) > 1:
            bar_width = (timestamps[-1] - timestamps[0]) / len(timestamps) * 0.8
        else:
            bar_width = 60  # 默认1分钟宽度
        
        # 使用ScatterPlotItem绘制柱状图
        for i, (t, vol) in enumerate(zip(timestamps, volumes)):
            # 根据涨跌决定柱状图颜色(收盘价与开盘价比较)
            is_up = chunk_data[i]['close'] >= chunk_data[i]['open']
            brush = red_brush if is_up else green_brush
            
            # 创建柱子
            bar = pg.PlotCurveItem(
                x=[t-bar_width/2, t-bar_width/2, t+bar_width/2, t+bar_width/2],
                y=[0, vol, vol, 0],
                fillLevel=0,
                brush=brush,
                pen=None
            )
            bars.append(bar)
            self.chart_widget.addItem(bar)
            
        return bars

    def load_all_data(self, data):
        """加载所有数据"""
        if not data:
            self.status_label.setText("没有可用的K线数据")
            return
            
        # 获取当前图表类型
        chart_type = self.type_combo.currentText()
        
        # 清除当前图表
        self.chart_widget.clear()
        
        # 重新添加十字光标和提示标签
        self.chart_widget.addItem(self.crosshair_v, ignoreBounds=True)
        self.chart_widget.addItem(self.crosshair_h, ignoreBounds=True)
        self.chart_widget.addItem(self.tooltip_label)
        
        if chart_type == "K线图":
            # 创建K线图对象
            item = CandlestickItem(data)
            self.chart_widget.addItem(item)
            
        elif chart_type == "趋势图":
            # 创建趋势线图表
            timestamps = [item['timestamp']/1000 for item in data]  # 转为秒
            closes = [item['close'] for item in data]
            
            # 使用蓝色线条表示趋势
            pen = pg.mkPen(color='b', width=2)
            line = pg.PlotDataItem(timestamps, closes, pen=pen)
            self.chart_widget.addItem(line)
            
        elif chart_type == "成交量":
            # 创建成交量柱状图
            self.create_volume_bars(data)
        
        # 自动调整视图以适应所有数据
        self.set_view_range_for_data(data)
    
    def set_view_range_for_data(self, data):
        """设置图表视图范围以适应所有数据"""
        if not data:
            return
            
        # 计算时间范围
        timestamps = [item['timestamp']/1000 for item in data]  # 转为秒
        x_min, x_max = min(timestamps), max(timestamps)
        
        # 获取当前图表类型
        chart_type = self.type_combo.currentText()
        
        # 根据图表类型设置不同的价格范围
        if chart_type == "K线图":
            # K线图包含高低点，需要考虑高低价格
            lows = [item['low'] for item in data]
            highs = [item['high'] for item in data]
            y_min, y_max = min(lows), max(highs)
            
            # 增加一些边距
            y_padding = (y_max - y_min) * 0.1
            y_min = y_min - y_padding
            y_max = y_max + y_padding
            
        elif chart_type == "趋势图":
            # 趋势图只显示收盘价
            closes = [item['close'] for item in data]
            y_min, y_max = min(closes), max(closes)
            
            # 增加一些边距
            y_padding = (y_max - y_min) * 0.1
            y_min = y_min - y_padding
            y_max = y_max + y_padding
            
        elif chart_type == "成交量":
            # 成交量图从0开始
            volumes = [item['volume'] for item in data]
            y_min, y_max = 0, max(volumes) * 1.1  # 增加10%边距
            
        # 设置图表的视图范围
        x_padding = (x_max - x_min) * 0.05  # 增加5%的时间边距
        x_min = x_min - x_padding
        x_max = x_max + x_padding
        
        # 设置视图范围
        self.chart_widget.setXRange(x_min, x_max)
        self.chart_widget.setYRange(y_min, y_max)

    def mouse_moved(self, pos):
        """处理鼠标移动事件"""
        if self.chart_widget.sceneBoundingRect().contains(pos):
            # 转换坐标
            mouse_point = self.chart_widget.plotItem.vb.mapSceneToView(pos)
            
            # 更新十字线位置
            self.crosshair_v.setPos(mouse_point.x())
            self.crosshair_h.setPos(mouse_point.y())
            
            # 更新位置信息
            timestamp = int(mouse_point.x())
            date_str = datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S")
            price = mouse_point.y()
            
            # 更新提示标签
            self.tooltip_label.setText(f"日期: {date_str}\n价格: {price:.4f}")
            self.tooltip_label.setPos(mouse_point.x(), mouse_point.y())
            
            # 更新底部状态栏
            self.position_label.setText(f"时间: {date_str} | 价格: {price:.4f}")
    
    def on_chart_type_changed(self, index):
        """图表类型变更处理"""
        chart_type = self.type_combo.currentText()
        self.status_label.setText(f"切换到 {chart_type}")
        
        # 加载所有数据
        all_data = ohlc_data_manager.get_all_data()
        if all_data:
            self.load_all_data(all_data)
        
    def on_zoom_changed(self, index):
        """缩放级别变更处理"""
        zoom_level = self.zoom_combo.currentText()
        all_data = ohlc_data_manager.get_all_data()
        
        if not all_data:
            return
            
        # 不同缩放级别
        if zoom_level == "全部":
            self.set_view_range_for_data(all_data)
        else:
            # 提取数字
            count = int(''.join(filter(str.isdigit, zoom_level)))
            
            # 使用最新的N条数据
            if len(all_data) > count:
                sub_data = all_data[-count:]
                self.set_view_range_for_data(sub_data)
            else:
                self.set_view_range_for_data(all_data)
    
    def on_view_changed(self, view_box):
        """处理视图变更事件"""
        # 不再需要加载数据块，仅用于记录当前视图状态
        pass
    
    def handle_viewport_change(self, direction, x_range):
        """处理视口变更"""
        # 不再需要加载数据块
        pass 